#!/usr/bin/python3
# -*- coding: utf-8; tab-width: 4; indent-tabs-mode: t -*-

import os
import logging
import subprocess
from gbs_util import GbsUtil
from gbs_param import GbsConst
from gbs_common import GbsSystem
from gbs_common import GbsPluginManager
from gbs_common import GbsProtocolException


class GbsCtrlSession:

    def __init__(self, param, user, clientIp, clientPort, clientPubKey):
        self.param = param

        self.sessionDir = os.path.join(GbsConst.runDir, "session-%s-%s-%s" % (user, clientIp, clientPort))
        os.makedirs(self.sessionDir, exist_ok=True)
        GbsUtil.chown(self.sessionDir, GbsConst.user, GbsConst.group)

        self.sysObj = GbsSystem(self.param, clientPubKey)

        self.pluginList = []
        self.stage = None

    def dispose(self):
        if hasattr(self, "stage"):
            if self.stage == "syncup":
                self._syncupStageEndHandler()
            elif self.stage == "working":
                self._workingStageEndHandler()
            else:
                assert self.stage is None
            del self.stage

        self._finiHandler()

        if hasattr(self, "pluginList"):
            del self.pluginList

        if hasattr(self, "sysObj"):
            self.sysObj.dispose()
            del self.sysObj

        if hasattr(self, "sessionDir"):
            GbsUtil.forceDelete(self.sessionDir)
            del self.sessionDir

    def onRequest(self, requestObj):
        if "command" not in requestObj:
            raise GbsProtocolException("Missing \"command\" in request object")

        if requestObj["command"] == "init":
            return self.cmdInit(requestObj)
        elif requestObj["command"] == "stage-syncup":
            return self.cmdStage("syncup", requestObj)
        elif requestObj["command"] == "stage-working":
            return self.cmdStage("working", requestObj)
        elif requestObj["command"] == "quit":
            return self.cmdQuit(requestObj)
        else:
            raise GbsProtocolException("Unknown command")

    def cmdInit(self, requestObj):
        logging.debug("Control Channel: Command \"init\" received from client \"%s\"." % (self._formatClient()))
        try:
            self._initHandler(requestObj)
            logging.debug("Control Channel: Command \"init\" processed from client \"%s\"." % (self._formatClient()))
            return {"return": {}}
        except BaseException as e:
            self._finiHandler()
            logging.exception("Control Channel: Command \"init\" error %s from client \"%s\"." % (str(e), self._formatClient()))
            return {"error": str(e)}

    def cmdStage(self, stage, requestObj):
        STAGE_LIST = ["syncup", "working"]
        assert stage in STAGE_LIST

        try:
            if STAGE_LIST.index(stage) == 0 and self.stage is not None:
                raise GbsProtocolException("Command \"stage-%s\" out of order" % (stage))
            if STAGE_LIST.index(stage) > 0 and STAGE_LIST.index(self.stage) + 1 != STAGE_LIST.index(stage):
                raise GbsProtocolException("Command \"stage-%s\" out of order" % (stage))

            # stage end processing
            if self.stage == "syncup":
                self._syncupStageEndHandler()
            elif self.stage == "working":
                self._workingStageEndHandler()
            else:
                assert self.stage is None

            # stage start processing
            logging.debug("Control Channel: Command \"stage-%s\" received from client \"%s\"." % (stage, self._formatClient()))
            self.stage = stage
            try:
                if self.stage == "syncup":
                    ret = self._syncupStageStartHandler(requestObj)
                elif self.stage == "working":
                    ret = self._workingStageStartHandler(requestObj)
                else:
                    assert False
                logging.debug("Control Channel: Command \"stage-%s\" processed from client \"%s\"." % (stage, self._formatClient()))
                return self._formatStageReturn(ret)
            except BaseException:
                if self.stage == "syncup":
                    self._syncupStageEndHandler()
                elif self.stage == "working":
                    self._workingStageEndHandler()
                else:
                    assert False
                self.stage = None
                raise
        except BaseException as e:
            logging.exception("Control Channel: Command \"stage-%s\" error %s from client \"%s\"." % (stage, str(e), self._formatClient()))
            return {"error": str(e)}

    def cmdQuit(self, requestObj):
        logging.debug("Control Channel: Command \"quit\" from client \"%s\"." % (self._formatClient()))
        self._finiHandler()
        return {"return": {}}

    def _initHandler(self, requestObj):
        if "hostname" in requestObj:
            logging.debug("Control Channel:     hostname = %s" % (requestObj["hostname"]))
            self.sysObj.getClientInfo().hostname = requestObj["hostname"]

        if "cpu-arch" in requestObj:
            logging.debug("Control Channel:     cpu-arch = %s" % (requestObj["cpu-arch"]))
            self.sysObj.getClientInfo().cpu_arch = requestObj["cpu-arch"]
        else:
            raise GbsProtocolException("Missing \"cpu-arch\" in command \"init\"")

        if "plugin" in requestObj:
            logging.debug("Control Channel:     plugin = %s" % (requestObj["plugin"]))
            for pname in requestObj["plugin"].split(","):
                self.pluginList.append(GbsPluginManager.loadPluginObject(pname, self.param, self))

        self.sysObj.mount()
        self.sysObj.commitClientInfo()

        for plugin in self.pluginList:
            if hasattr(plugin, "init_handler"):
                plugin.init_handler(requestObj)

    def _finiHandler(self):
        # should raise no exception
        for plugin in reversed(self.pluginList):
            if hasattr(plugin, "fini_handler"):
                plugin.fini_handler()
        self.sysObj.unmount()

    def _syncupStageStartHandler(self, requestObj):
        ret = {}

        # invoke plugin start handler
        for plugin in self.pluginList:
            if hasattr(plugin, "stage_syncup_start_handler"):
                ret2 = plugin.stage_syncup_start_handler(requestObj)
                GbsUtil.mergeDictWithOverwriteAsException(ret, ret2)

        # start rsync service
        self.rsyncServ = _RsyncService(self.param, self.sysObj.uuid, self.sslSock.getpeername()[0],
                                       self.sslSock.get_peer_certificate(), self.sysObj.mntDir, True)
        self.rsyncServ.start()
        logging.debug("Control Channel:     rsync service on %d" % (self.rsyncServ.getPort()))
        GbsUtil.mergeDictWithOverwriteAsException(ret, {
            "rsync-port": self.rsyncServ.getPort(),
        })

        return ret

    def _syncupStageEndHandler(self):
        # should raise no exception
        if hasattr(self, "rsyncServ"):
            self.rsyncServ.stop()
            del self.rsyncServ
        for plugin in reversed(self.pluginList):
            if hasattr(plugin, "stage_syncup_end_handler"):
                plugin.stage_syncup_end_handler()

    def _workingStageStartHandler(self, requestObj):
        ret = {}

        # prepare root directory
        self.sysObj.prepareRoot()

        # invoke plugin start handler
        for plugin in self.pluginList:
            if hasattr(plugin, "stage_working_start_handler"):
                ret2 = plugin.stage_working_start_handler(requestObj)
                GbsUtil.mergeDictWithOverwriteAsException(ret, ret2)

        # start rsync service
        self.rsyncServ = RsyncService(self.param, self.sysObj.uuid, self.sysObj.mntDir, False)
        self.rsyncServ.start()
        logging.debug("Control Channel:     rsync service on %d" % (self.rsyncServ.getPort()))
        GbsUtil.mergeDictWithOverwriteAsException(ret, {
            "rsync-port": self.rsyncServ.getPort(),
        })

        return ret

    def _workingStageEndHandler(self):
        # should raise no exception
        if hasattr(self, "rsyncServ"):
            self.rsyncServ.stop()
            del self.rsyncServ
        for plugin in reversed(self.pluginList):
            if hasattr(plugin, "stage_working_end_handler"):
                plugin.stage_working_end_handler()
        self.sysObj.unPrepareRoot()

    def _formatStageReturn(self, ret):
        ret["stage"] = self.stage
        return {"return": ret}

    def _formatClient(self):
        ret = "UUID:%s" % (self.sysObj.uuid)
        if self.sysObj.getClientInfo().hostname is not None:
            ret = "%s(%s)" % (self.sysObj.getClientInfo().hostname, ret)
        return ret


class GbsSessionException(Exception):
    pass


class GbsPluginException(Exception):
    pass


class _RsyncService:

    def __init__(self, param, uuid, rootDir, upOrDown):
        self.param = param
        self.uuid = uuid
        self.rootDir = rootDir
        self.upOrDown = upOrDown

        self.rsyncdCfgFile = os.path.join(self.param.tmpDir, uuid + "-rsyncd.conf")
        self.rsyncdLockFile = os.path.join(self.param.tmpDir, uuid + "-rsyncd.lock")
        self.rsyncdLogFile = os.path.join(self.param.logDir, uuid + "-rsyncd.log")

        self.rsyncPort = None

        self.rsyncProc = None

    def start(self):
        try:
            self.rsyncPort = GbsUtil.getFreeTcpPort()
            self.rsyncProc = self._runRsyncDeamon()
            GbsUtil.waitTcpPort(self.rsyncPort)
        except:
            self.stop()
            raise

    def stop(self):
        if self.rsyncProc is not None:
            self.rsyncProc.terminate()
            self.rsyncProc.wait()
        GbsUtil.forceDelete(self.rsyncdLockFile)
        GbsUtil.forceDelete(self.rsyncdCfgFile)

    def getPort(self):
        return self.rsyncPort

    def _runRsyncDeamon(self):
        buf = ""
        buf += "lock file = %s\n" % (self.rsyncdLockFile)
        buf += "log file = %s\n" % (self.rsyncdLogFile)
        buf += "\n"
        buf += "port = %s\n" % (self.rsyncPort)
        buf += "max connections = 1\n"
        buf += "timeout = 600\n"
        buf += "hosts allow = 127.0.0.1\n"
        buf += "\n"
        buf += "use chroot = yes\n"
        buf += "uid = root\n"
        buf += "gid = root\n"
        buf += "\n"
        buf += "[main]\n"
        buf += "path = %s\n" % (self.rootDir)
        buf += "read only = %s\n" % ("no" if self.upOrDown else "yes")
        with open(self.rsyncdCfgFile, "w") as f:
            f.write(buf)

        cmd = ""
        cmd += "rsync --daemon --no-detach --config=\"%s\"" % (self.rsyncdCfgFile)
        proc = subprocess.Popen(cmd, shell=True, universal_newlines=True)
        return proc
